#!/usr/bin/env python
# coding=utf-8

from __future__ import print_function
import json
import easydict
import argparse
import os.path
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import time
from skimage import io
from sort.sort_impl import Sort

display = False
config = easydict.EasyDict()
config.DEBUG = False

# 通过躯干box 跟踪, 再估计整体位置
I_MODE_JOINT = 1
# pose boxes
I_MODE_POSE_BOX = 2
# 基于上一个, 修复宽度问题
I_MODE_FIX_WIDTH = 3

MODE = I_MODE_FIX_WIDTH
# MODE = -1 # don t fix

# L_SEQUENCE_NAMES = ['ADL-Rundle-1',
#                     'AVG-TownCentre',
#                     'ETH-Jelmoli',
#                     'KITTI-16',
#                     'PETS09-S2L2',
#                     'Venice-1',
#                     'ADL-Rundle-3',
#                     'ETH-Crossing',
#                     'ETH-Linthescher',
#                     'KITTI-19',
#                     'TUD-Crossing']
L_SEQUENCE_NAMES = ['ADL-Rundle-6',
            'ADL-Rundle-8',
            'ETH-Bahnhof',
            'ETH-Pedcross2',
            'ETH-Sunnyday',
            'KITTI-13',
            'KITTI-17',
            'PETS09-S2L1',
            'TUD-Campus',
            'TUD-Stadtmitte',
            'Venice-2']
set_name = 'train_v1'
# set_name = 'test'

output_dir_name = '{}_mode_{}_result'.format(set_name, MODE)


def tlwh2():
    pass


def estimate_box(tlwh_joint_box):
    b1 = tlwh_joint_box
    dw, dh = b1[2], b1[3]
    scale_w, scale_h = (0.4, 0.31)
    pos_shift = (-130 * (dw / 530), -200 * (dh / 1080.))
    new_box = np.zeros_like(b1)
    new_box[0] = pos_shift[0] + b1[0]
    new_box[1] = pos_shift[1] + b1[1]
    new_box[2] = dw / scale_w
    new_box[3] = dh / scale_h
    return new_box


class BoxEstimator:
    def __init__(self, param_file):
        param = np.loadtxt(param_file)
        self.scale_w = param[0]
        self.scale_h = param[1]
        self.offset = param[2]

    def estimate_box2(self, tlwh_joint_box, scale_w, scale_h, v_offset_ratio):
        ib = tlwh_joint_box
        dw, dh = ib[2], ib[3]
        # scale_w, scale_h = (0.4, 0.31)
        # cx = ib[0] + ib[2]/2
        # cy = ib[1] + ib[3]/2
        new_box = np.zeros_like(ib)
        # new_cy = cy + shift
        new_box[2] = dw / scale_w
        new_box[3] = dh / scale_h
        ip = ib[[0, 1]] + ib[[2, 3]] / 2.
        op_v = ip[1] - v_offset_ratio * dh
        new_box[0] = ip[0] - new_box[2] / 2.
        new_box[1] = op_v - new_box[3] / 2.
        new_box[0] = 0 if new_box[0] < 0 else new_box[0]
        new_box[1] = 0 if new_box[1] < 0 else new_box[1]
        return new_box

    def estimate(self, box):
        return self.estimate_box2(box, self.scale_w, self.scale_h, self.offset)


def oneshot_origin(det_path, output_path):
    # det_path = '/Users/gerrie/prj/end2end_mot/output/mymot_det_test.txt'
    mot_tracker = Sort()
    # seq_dets = np.loadtxt(det_path.format(seq_name), delimiter=',')  # load detections
    seq_dets = np.loadtxt(det_path, delimiter=',')
    with open(output_path, 'w') as out_file:
        t1 = time.time()
        for frame in range(int(seq_dets[:, 0].max())):
            # frame += 1  # detection and frame numbers begin at 1
            dets = seq_dets[seq_dets[:, 0] == frame, 2:7]
            dets[:, 2:4] += dets[:, 0:2]  # convert to [x1,y1,w,h] to [x1,y1,x2,y2]
            trackers = mot_tracker.update(dets)
            for d in trackers:
                tid = d[4]
                dd = d[0], d[1], d[2] - d[0], d[3] - d[1]
                # dd = estimate_box(tlwh_box)
                print('%d,%d,%.2f,%.2f,%.2f,%.2f,1,-1,-1,-1' % (frame, tid, dd[0], dd[1], dd[2], dd[3]),
                      file=out_file)
        t2 = time.time()
        print('fps : {}'.format(seq_dets[:, 0].max()/(t2 - t1)))

import pandas as pd


def oneshot_df():
    det_path = '/Users/gerrie/prj/end2end_mot/output/mymot_det_test.txt'
    mot_tracker = Sort()
    df = pd.read_csv()


from sort.util import fix_box_width

def oneshot_15(det_path, output_path):
    """

    :param det_path:
    :param output_path:
    :return:
    """
    # det_path = '/Users/gerrie/prj/end2end_mot/deep-sort/MOT16/train_v1/{}/det/det.txt'
    # det_path = '/Users/gerrie/prj/end2end_mot/tmp/det-{}.txt'
    # det_path = '/Users/gerrie/prj/end2end_mot/output/mymot_joints_det_{}.txt'
    # seq_dets = np.loadtxt(det_path.format(seq_name),delimiter=',') #load detections

    mot_tracker = Sort(iou_threshold=0.3)
    # mot_tracker = Sort(max_age=7, min_hits=2, iou_threshold=0.2)
    seq_dets = np.loadtxt(det_path, delimiter=',')  # load detections

    estimator = BoxEstimator('../output/point_box_param.txt')

    with open(output_path, 'w') as out_file:
        print("dealing with: " + output_path + "...")
        t1 = time.time()
        for frame in range(int(seq_dets[:, 0].max())):
            frame += 1  # detection and frame numbers begin at 1
            dets = seq_dets[seq_dets[:, 0] == frame, 2:7]
            dets[:, 2:4] += dets[:, 0:2]  # convert to [x1,y1,w,h] to [x1,y1,x2,y2]
            trackers = mot_tracker.update(dets)
            for d in trackers:
                tlwh_box = [d[0], d[1], d[2] - d[0], d[3] - d[1]]
                tid = d[4]
                # dd = estimate_box(tlwh_box)
                dd = tlwh_box
                if MODE == I_MODE_JOINT:
                    dd = estimator.estimate(np.array(tlwh_box))
                elif MODE == I_MODE_FIX_WIDTH:
                    dd = fix_box_width(tlwh_box)

                print('%d,%d,%.2f,%.2f,%.2f,%.2f,1,-1,-1,-1' % (frame, tid, dd[0], dd[1], dd[2], dd[3]),
                      file=out_file)
        t2 = time.time()
        print('fps : {}'.format(seq_dets[:, 0].max()/(t2 - t1)))


def oneshot_15_pose(pose_df_path, output_path):
    import pandas as pd
    mot_tracker = Sort(iou_threshold=0.1)
    df = pd.read_hdf(pose_df_path, 'df')
    with open(output_path, 'w') as out_file:
        for frame in np.unique(df.fid.values):
            # frame += 1
            print(frame)
            cur_df = df[df.fid == frame]
            ser_out_boxes = cur_df.box
            # ser_out_boxes.iloc[0]
            # print(dets.shape)


def mot15_batch_origin():
    _set_name = 'train'
    for seq in L_SEQUENCE_NAMES:
        det_path = os.path.join('../MOT15', _set_name, seq, 'det/det.txt')
        det_path = os.path.join('../FrCNN_Det', seq, 'det.txt')
        output_path = os.path.join('../output', 'origin', _set_name)
        if not os.path.exists(output_path):
            os.makedirs(output_path)
        output_path = os.path.join(output_path, seq + '.txt')
        oneshot_origin(det_path, output_path)


def mot15_batch_handle():
    det_path = ''
    output = '../output/{}/{}.txt'
    if MODE == I_MODE_JOINT:
        det_path = '../output/mymot_joints_det_{}.txt'  #
        # output_dir_name = 'train_result'
    elif MODE == I_MODE_POSE_BOX or MODE == I_MODE_FIX_WIDTH:
        # det_path = '../output/mot15_pose_acc/{}.txt'
        det_path = '../output/mot15_pose_det/' + set_name + '/{}.txt'
        # output_dir_name = 'mot15_pose_result'
    for seq in L_SEQUENCE_NAMES:
        if not os.path.exists(det_path.format(seq)):
            print('miss file ' + det_path.format(seq))
            continue
        if not os.path.exists('../output/{}'.format(output_dir_name)):
            os.makedirs('../output/{}'.format(output_dir_name))
        oneshot_15(det_path.format(seq), output.format(output_dir_name, seq))


def mot15_batch_df():
    seq = 'TUD-Stadtmitte'
    df_path = '../output/mot15_pose/train/{}/det_pose.hdf5'.format(seq)

    # output = '../output/{}/{}.txt'
    output_path = 'tmp.txt'
    oneshot_15_pose(df_path, output_path)


if __name__ == '__main__':
    # seq_name = 'MOT16-02'
    # seq_name = 'PETS09-S2L1'
    # output_path = '../output/train_result/{}.txt'.format(seq_name)
    # oneshot_16(seq_name, output_path)

    # mot15_batch_df()
    # mot15_batch_handle()
    mot15_batch_origin() #原始版本 sort作者自制det数据

    # det_path = '../output/mymot_joints_det_{}.txt'
    # oneshot_15(det_path.format(seq_name), output_path)





#
# def oneshot_15_pose_bak(det_path, det_path_b, output_path):
#     mot_tracker = Sort(iou_threshold=0.1)
#     seq_dets = np.loadtxt(det_path, delimiter=',')  # load detections
#     seq_dets_b = np.loadtxt(det_path_b, delimiter=',')  # load detections
#
#     with open(output_path, 'w') as out_file:
#         for frame in range(int(seq_dets[:, 0].max())):
#             frame += 1  # detection and frame numbers begin at 1
#             dets = seq_dets[seq_dets[:, 0] == frame, 2:7]
#             dets[:, 2:4] += dets[:, 0:2]  # convert to [x1,y1,w,h] to [x1,y1,x2,y2]
#             trackers = mot_tracker.update(dets)
#             for d in trackers:
#                 tlwh_box = [d[0], d[1], d[2] - d[0], d[3] - d[1]]
#                 tid = d[4]
#                 # dd = estimate_box(tlwh_box)
#                 dd = tlwh_box
#                 print('%d,%d,%.2f,%.2f,%.2f,%.2f,1,-1,-1,-1' % (frame, tid, dd[0], dd[1], dd[2], dd[3]),
#                       file=out_file)
